from langchain_core.embeddings import Embeddings
from sentence_transformers import SentenceTransformer

class BGE_Embed(Embeddings):
    # def __init__(self, model_path: str = "/data/ai/langchain/yunjian/model/AI-ModelScope/bge-large-zh-v1.5"):
    def __init__(self, model_path: str = "/root/.cache/modelscope/hub/AI-ModelScope/bge-large-zh-v1___5"):

        self.embedding = SentenceTransformer(model_path)

    # 对文档进行词嵌入
    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        return [self.embed_query(text) for text in texts]

    # 对查询语句进行嵌入
    def embed_query(self, text: str) -> list[float]:
        return self.embedding.encode(text).tolist()

# 示例使用
if __name__ == "__main__":
    embedder = BGE_Embed()

    # 示例文档
    documents = ["人工智能的发展趋势", "机器学习的应用"]
    embeddings = embedder.embed_documents(documents)
    print("文档嵌入向量:", len(embeddings))

    # 示例查询
    query = "深度学习的进步"
    query_embedding = embedder.embed_query(query)
    print("查询嵌入向量:", query_embedding)



